import json
from funasr import AutoModel
import gradio as gr
import json
from docx import Document
from docx.shared import Pt
import colorsys

# 自定义CSS，用于将音频播放器固定在页面底部，并使识别输出文本占满整个页面
custom_css = """
        .speaker {
            height: 32px;
            display: flex;
            align-items: center;
            pointer-events: auto;
            border-radius: 4px;
            cursor: pointer;
            user-select: none;
        }
        .speaker-avatar {
            position: relative;
            top: 6px;
            float: left;
            height: 30px;
            width: 30px;
            border-radius: 16px;
            margin: 0 16px;
            background-color: #d7d7d7;
            text-align: center;
            line-height: 30px;
            font-size: 12px;
            color: #fff;
        }

        .speaker-name {
            margin-right: 10px;
            height: 18px;
            line-height: 18px;
            font-size: 14px;
            font-weight: 500;
            font-family: PingFangSC-Medium, PingFang SC;
            color: rgba(0, 0, 0, .66);
            max-width: 15em;
            white-space: nowrap;
            text-overflow: ellipsis;
            -webkit-user-select: none;
            -moz-user-select: none;
            -ms-user-select: none;
            user-select: none;
            overflow: hidden;
        }

        .paragraph-head {
            user-select: none;
            -webkit-user-select: none;
            -moz-user-select: none;
            -ms-user-select: none;
            pointer-events: none;
            height: 18px;
            display: flex;
            align-items: center;
        }

        .seek-bar {
            padding: 12px 0;
            box-sizing: border-box;
        }
        .editor {
            outline: none;
        }
        .paragraph {
            word-break: break-word;
            user-select: text;
            padding: 0 10px 20px 53px;
            line-height: 22px;
            transform: translateZ(0);
            outline: 0;
            margin: 5px 0 0 10px;
            font-size: 14px;
        }
"""
base_color = [0, 169, 152]

Contentformat = {
    "html": {
        "bgColors": [
            "background-color: rgb(255, 137, 82);",
            "background-color: rgb(0, 169, 152);",
            "background-color: rgb(191, 142, 238);",
            "background-color: rgb(255, 88, 90);",
            "background-color: rgb(148, 63, 230);",
            "background-color: rgb(63, 191, 90);"
        ],
        "paragraph": "<p class='paragraph active' style='direction: ltr;'>{paragraph}</p>",
        "resultTxt": "<div class='paragraph-head'>"
                     "<div class='speaker el-tooltip__trigger el-tooltip__trigger'> \
                     <div class='speaker-avatar' style='{bgColor}'>{currentSpk}</div> \
                     <div class='speaker-name'>说话人{currentSpk}</div> \
                     <div class='time-code' style='padding-left: 0px;'>{startTime}~{endTime}</div> \
                     </div> \
                     </div> \
                     {text}"
    },
    "docx": {
        "bgColors": [
            "background-color: rgb(255, 137, 82);",
            "background-color: rgb(0, 169, 152);",
            "background-color: rgb(191, 142, 238);",
            "background-color: rgb(255, 88, 90);",
            "background-color: rgb(148, 63, 230);",
            "background-color: rgb(63, 191, 90);"
        ],
        "paragraph": "{paragraph}\n",
        "resultTxt": "\n说话人{currentSpk} {startTime}~{endTime}\n{text}"
    }
}

fileResult = "temp_audio_file.json"
exportFile = "output.docx"


def time2hour(timestamp):
    # 如果时间戳小于3600秒，将其转换为分钟和秒
    if timestamp < 3600:
        hours = 0
        minutes = timestamp // 60
        seconds = timestamp % 60
    else:
        hours = timestamp // 3600
        minutes = (timestamp % 3600) // 60
        seconds = (timestamp % 3600) % 60

    # 格式化为 "小时:分钟:秒"
    formattedTime = f"{hours:02}:{minutes:02}:{seconds:02}"
    return formattedTime


def recognizeAudio(audio):
    data = ""
    fileResult = voice2json(audio)
    with open(fileResult, 'r', encoding='utf-8') as file:
        data = json.load(file)
    sentenceInfo = data[0]["sentence_info"]
    return "<div class='editor'>{}</div>".format(parseContent(sentenceInfo, "html"))


def parseContent(sentenceInfo, docType):
    text = ""
    resultTxt = ""
    paragraph = ""
    paragraphEnd = ["。", "？", "！"]
    currentSpk = sentenceInfo[0]["spk"]
    startTime = time2hour(int(sentenceInfo[0]["start"] / 1000))
    endTime = time2hour(int(sentenceInfo[0]["end"] / 1000))
    spks = len(set([item["spk"] for item in sentenceInfo]))
    bgColors = generate_colors(base_color, spks)
    # bgColors = Contentformat[docType]["bgColors"]
    paragraphFormat = Contentformat[docType]["paragraph"]
    resultTxtFormat = Contentformat[docType]["resultTxt"]

    for info in sentenceInfo:
        spk = info["spk"]
        if spk == currentSpk:
            paragraph += info['text']
            if info['text'][-1] in paragraphEnd:
                text += paragraphFormat.format(paragraph=paragraph)
                paragraph = ""
            endTime = time2hour(int(info["end"] / 1000))
        else:
            text += paragraphFormat.format(paragraph=paragraph)
            bgColor = bgColors[currentSpk - 1]
            resultTxt = resultTxt + resultTxtFormat.format(bgColor=bgColor, currentSpk=currentSpk, startTime=startTime,
                                                           endTime=endTime, text=text)
            paragraph = info['text']
            startTime = time2hour(int(info["start"] / 1000))
            endTime = time2hour(int(info["end"] / 1000))
            currentSpk = spk
            text = ""

    info_end = sentenceInfo[-1]
    endTime = time2hour(int(info_end["end"] / 1000))
    bgColor = bgColors[currentSpk - 1]
    resultTxt = resultTxt + resultTxtFormat.format(bgColor=bgColor, currentSpk=currentSpk, startTime=startTime,
                                                   endTime=endTime, text=text)
    return resultTxt


def add_paragraph_with_style(document, text, style):
    p = document.add_paragraph()
    run = p.add_run(text)
    run.font.size = style.get('font_size', Pt(12))
    run.font.name = '宋体'
    return p


def generateFile():
    globalStyle = {
        'font_size': Pt(12),
        'font.name': '宋体'
    }

    with open(fileResult, 'r', encoding='utf-8') as file:
        data = json.load(file)

    sentenceInfo = data[0]["sentence_info"]
    docx_content = parseContent(sentenceInfo, "docx")
    doc = Document()
    add_paragraph_with_style(doc, docx_content, globalStyle)

    # 保存 DOCX 文件
    doc.save(exportFile)

    return exportFile


def voice2json(fileVoice):
    print("starting")
    model = AutoModel(model="paraformer-zh", model_revision="v2.0.4",
                      vad_model="fsmn-vad", vad_model_revision="v2.0.4",
                      punc_model="ct-punc-c", punc_model_revision="v2.0.4",
                      spk_model="cam++", spk_model_revision="v2.0.2")
    res = model.generate(input=fileVoice,
                         batch_size_s=300,
                         hotword='魔搭')
    json_string = json.dumps(res, ensure_ascii=False)
    with open(fileResult, 'w', encoding='utf-8') as file:
        file.write(json_string)
    print("end")

    return fileResult


def generate_colors(base_rgb, num_colors):
    # 将RGB值标准化到0-1范围
    base_rgb = [x / 255.0 for x in base_rgb]
    # 将RGB转换为HSV
    base_hsv = colorsys.rgb_to_hsv(*base_rgb)

    colors = []
    for i in range(num_colors):
        # 旋转色相
        new_hue = (base_hsv[0] + i / num_colors) % 1.0
        # 保持饱和度和亮度不变，生成新的颜色
        new_rgb = colorsys.hsv_to_rgb(new_hue, base_hsv[1], base_hsv[2])
        # 转换回0-255范围的RGB
        new_rgb = [int(x * 255) for x in new_rgb]
        colors.append("background-color: rgb({}, {}, {});".format(new_rgb[0], new_rgb[1], new_rgb[2]))

    return colors


# css=custom_css
with gr.Blocks(css=custom_css) as demo:
    with gr.Row(visible=True) as inputRow:
        with gr.Column(scale=1):
            audioInput = gr.Audio(type="filepath", label="音频文件", sources="upload")

    with gr.Row(visible=True) as audioPlayerRow:
        recognizeButton = gr.Button("开始识别")

    with gr.Row(visible=False) as audioPlayerButton:
        resetButton = gr.Button("重新识别", min_width=10)
        exportButton = gr.DownloadButton("下载结果", min_width=10, value=exportFile)

    with gr.Row(visible=False) as resultRow:
        with gr.Column(scale=3):
            recognizedText = gr.HTML()


    def updateInterface(audioPath):
        text = recognizeAudio(audioPath)
        htmlContent = "<div style='height: 75vh;overflow-y: auto;scrollbar-width: none;'>{}</div>".format(
            text.replace('\n', '<br>'))
        return gr.update(visible=True), gr.update(visible=False), gr.update(visible=True), gr.update(
            visible=True), htmlContent


    def resetPage():
        return gr.update(visible=True), gr.update(visible=True), gr.update(visible=False), gr.update(visible=False)


    recognizeButton.click(
        fn=updateInterface,
        inputs=audioInput,
        outputs=[inputRow, audioPlayerRow, audioPlayerButton, resultRow, recognizedText]
    )

    resetButton.click(
        fn=resetPage,
        outputs=[inputRow, audioPlayerRow, audioPlayerButton, resultRow])

if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0")
